from Yolo.darknet import Darknet19
from DataSet.TongueDS import TongueDS


import numpy as np
from torch.autograd import Variable
from torch.utils.data import DataLoader

import torch.nn as nn

model = Darknet19().cuda()
model = nn.DataParallel(model)

DS = TongueDS()
DL = DataLoader(DS,batch_size=4,shuffle=True)

for i,out in enumerate(DL):
    break

def DLTransform(out):
    images = Variable(out['images'].permute(0,3,1,2)).float().cuda()
    gt_boxes = out['gt_boxes']
    boxes_list = []
    for item in gt_boxes:
        boxes_list.append(item.numpy()[np.newaxis,:])
    gt_classes = out['gt_classes']
    classes_list = []
    for item in gt_classes:
        classes_list.append(item.numpy())
    empty_list = [[] for _ in range(len(boxes_list))]
    return images,boxes_list,classes_list,empty_list,0

out = DLTransform(out)

len(out)
out[-2]

output = model(*out)


len(output)

output[0].shape
output[1].shape
output[2].shape

model.loss
